import os
import numpy as np
import cv2
import re
from concurrent.futures import ThreadPoolExecutor


def extract_prefix(filename):
    pattern = re.compile(r'(.*)tr\d+-tc\d+\.png')
    match = pattern.match(filename)
    if match:
        return match.group(1)
    return ""


def cut_and_save_large_image(image_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    image_files = sorted([f for f in os.listdir(image_dir) if f.endswith('.png')])

    def process_image(image_file):
        prefix = extract_prefix(image_file)
        # 解析行和列编号
        match = re.search(r'tr(\d+)-tc(\d+)', image_file)
        if not match:
            return

        big_tr = int(match.group(1))
        big_tc = int(match.group(2))

        image_path = os.path.join(image_dir, image_file)
        large_image = cv2.imread(image_path)
        if large_image is None:
            raise FileNotFoundError(f"Unable to open image file: {image_path}")

        large_tile_size = 4096

        row_start = 1 + 5 * (big_tr - 1)
        row_end = row_start + 6
        col_start = 1 + 5 * (big_tc - 1)
        col_end = col_start + 6
        if big_tr == 1:
            row_valid_start = 1
            row_valid_end = 6
        elif big_tr == 7:
            row_valid_start = 32
            row_valid_end = 37
        else:
            row_valid_start = 2 + (big_tr - 1) * 5
            row_valid_end = row_valid_start + 4
        if big_tc == 1:
            col_valid_start = 1
            col_valid_end = 6
        elif big_tc == 7:
            col_valid_start = 36
            col_valid_end = 37
        else:
            col_valid_start = 2 + (big_tc - 1) * 5
            col_valid_end = col_valid_start + 4

        for small_row in range(row_start, row_end + 1):
            if small_row < row_valid_start or small_row > row_valid_end:
                continue
            for small_col in range(col_start, col_end + 1):
                if small_col < col_valid_start or small_col > col_valid_end:
                    continue

                x_offset = (small_row - row_start) * large_tile_size
                y_offset = (small_col - col_start) * large_tile_size

                small_image = large_image[x_offset:x_offset + large_tile_size, y_offset:y_offset + large_tile_size]

                small_filename = f"{prefix}tr{small_row}-tc{small_col}.png"
                small_image_path = os.path.join(output_dir, small_filename)
                cv2.imwrite(small_image_path, small_image)
                print(f"Saved {small_filename} to {small_image_path}")

    with ThreadPoolExecutor(max_workers=7) as executor:
        executor.map(process_image, image_files)


# 示例用法
image_dir = '/root/data/deep_align/11867_32nm_to_8nm_fine_big_align'  # 替换为实际的大图文件夹路径
output_dir = '/root/data/deep_align/11867_32nm_to_8nm_fine_align'  # 替换为实际的小图输出文件夹路径

cut_and_save_large_image(image_dir, output_dir)
